import argparse
import socket
import subprocess
import json
import os
import time


def connect_to_child(host, port):
    """连接到子程序服务端并发送数据"""
    init_params="Hello from Python3 main program"
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(1000)
        sock.connect((host, port))

        # 发送初始化参数
        header = {
            "auth": {"username": "admin", "password": "123"},
            "type": "normal",
            "data": init_params
        }
        sock.sendall(json.dumps(header).encode())
        response = sock.recv(4096)
        print(f"[主程序] 收到响应: {response.decode()}")
        sock.close()
    except Exception as e:
        print(f"连接错误: {str(e)}")


def stream_generator(host, port, init_params=10):
    sock = socket.socket()
    sock.settimeout(1000)
    sock.connect((host, port))

    # 发送初始化参数
    header = {
        "auth": {"username": "admin", "password": "123"},
        "type": "stream",
        "data": init_params
    }
    sock.sendall(json.dumps(header).encode())

    while True:
        data = sock.recv(4096)
        if data == b"STREAM_END":
            break
        yield data.decode('utf-8')
    sock.close()


def start_local_child():
    """启动本地Python2子程序并获取端口号"""
    explain = r"D:/SIMULIA/EstProducts/2020/win_b64/tools/SMApy/python2.7/python.exe "
    script_path = r"G:\projects\python-projicts\abq-parameter-analysis-framework\examples\test\py2_socket_test.py"
    # 启动子程序并捕获其输出
    proc = subprocess.Popen(
        f"cmd.exe /C {explain}{script_path}",
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        text=True,
        cwd=r"" + os.getcwd()
    )

    # 获取子程序的进程ID
    pid = proc.pid
    proc.stdin.write(json.dumps({"pid": pid}) + "\n")
    proc.stdin.flush()
    # 构建文件名
    filename = f"socket_info_process_{pid}.json"
    print(f"[主程序] 需求文件名: {filename}")

    # 等待文件生成并读取端口号
    while not os.path.exists(filename):
        time.sleep(0.05)

    # 读取文件内容
    with open(filename, "r") as f:
        socket_info = json.load(f)
        port = socket_info["port"]

    # 删除文件
    os.remove(filename)
    print(f"[主程序] 子程序端口: {port}")
    return port


def test_common(port=None):
    if port is None:
        port = start_local_child()
    connect_to_child("localhost", port)
    pass


def test_stream(port=None):
    if port is None:
        port = start_local_child()
    for data in stream_generator("localhost", port):
        print(data)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", help="远程子程序地址")
    parser.add_argument("--port", type=int, help="远程子程序端口")
    args = parser.parse_args()

    if args.host or args.port:
        # 连接指定远程子程序
        if not (args.host and args.port):
            print("必须同时提供 --host 和 --port 参数")
            return
        connect_to_child(args.host, args.port)
    else:
        # 自动启动本地子程序
        port = start_local_child()
        connect_to_child("localhost", port)


if __name__ == "__main__":
    # main()
    port=None
    test_stream(port)
    # test_common(port)
